# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-03-31 15:10:31
# @Last Modified by:   bao
# @Last Modified time: 2017-05-01 14:40:10
import sys
sys.path.append("..")

import re
import random
from CoreSim.MScheduler import *
from CoreSim.Flow import *
from CoreSim.Event import *
from CoreSim.Link import *
from CoreSim.Unit import *

Scatter = 0
Ring = 1

def assignNewFlow(srcHost, dstHost, flows, flowSize, startTime, mcastFlowId, numHosts, numToRs, numCores, subflowId =0):
    # unicast flow
    flow = Flow()
    flow.startId = srcHost
    flow.endId = dstHost
    flow.SetFlowSize(flowSize)
    flow.startTime = startTime
    flow.flowId = len(flows) + 1
    flow.mcastFlowId = mcastFlowId
    flow.subflowId = subflowId
    flows.append(flow)
    return flow

class Host(object):
    """docstring for Host"""
    def __init__(self):
        self.hostId = 0
        self.rSubFlowIDs = []
        self.sSubFlowIDs = []
        self.rBusy = False
        self.sBusy = False
        self.subFlowSet = []

    def __repr__(self):
        return 'hostId:{!r}, rSubFlowIDs:{!r}, sSubFlowIDs:{!r}, rBusy:{!r}, sBusy:{!r}, subFlowSet:{!r}'.format\
        (self.hostId, self.rSubFlowIDs, self.sSubFlowIDs, self.rBusy, self.sBusy, self.subFlowSet)

class MSchedulerRing(MScheduler):
    '''
    Multicast scheduler
    '''
    def __init__(self, topo_args, flow_args, conf_args):
        MScheduler.__init__(self, topo_args, flow_args, conf_args)
        '''
        For each node, maitains the received subflows and the on going sending subflows
        '''
        self.mcastHosts = {} # mcastflowId:{hostid:Host}
        self.mcastSubtree = {} # binomial tree
        self.mcastStage = {} # mcastflowId: stage
        self.type = 'ring'
    
    def InitialMulticast(self, eventQueue):
        '''
        Split the data based on binomial tree (scatter)
        '''
        for mflowId in self.mcastFlows:
            self.mcastStage[mflowId] = Scatter
            mcastFlow = self.mcastFlows[mflowId]
            
            index = mcastFlow.groupSize/2
            srcHost = mcastFlow.srcHost
            dstHost = mcastFlow.dstHosts.keys()[index]

            self.mcastSubtree[mflowId] = {}
            self.mcastSubtree[mflowId][srcHost] = mcastFlow.dstHosts.keys()[:index]
            self.mcastSubtree[mflowId][dstHost] = mcastFlow.dstHosts.keys()[index+1:]
            #print srcHost, dstHost
            #print mcastFlow.dstHosts.keys()
            #print 'mcastSubtree', self.mcastSubtree
            
            flow = assignNewFlow(srcHost, dstHost, self.flows, mcastFlow.flowSize/2, mcastFlow.startTime, \
                mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

            # eventqueue
            event = FlowArrivalEvent(mcastFlow.startTime, flow)
            eventQueue.push(event, mcastFlow.startTime)
    
    def UpdateMulticast(self, mflowId, curFlow, curTime, eventQueue):
        mcastFlow = self.mcastFlows[mflowId]
        size = mcastFlow.groupSize+1
        sub_flowSize = mcastFlow.flowSize/size

        if self.mcastStage[mflowId] == Scatter:
            mcastFlow.dstHosts[curFlow.endId] = True
            if not all(mcastFlow.dstHosts.values()):
                srcHost = curFlow.startId
                dstHost = curFlow.endId
                srcSubHosts = self.mcastSubtree[mflowId][srcHost]
                dstSubHosts = self.mcastSubtree[mflowId][dstHost]

                if srcSubHosts:
                    index1 = len(srcSubHosts)/2
                    host1 = srcSubHosts[index1]
                    self.mcastSubtree[mflowId][srcHost] = srcSubHosts[:index1]
                    self.mcastSubtree[mflowId][host1] = srcSubHosts[index1+1:]

                    flow = assignNewFlow(srcHost, host1, self.flows, curFlow.flowSize/2, curTime, \
                                        mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)
                    
                    event = FlowArrivalEvent(curTime, flow)
                    eventQueue.push(event, event.time)
                
                if dstSubHosts:
                    index2 = len(dstSubHosts)/2
                    host2 = dstSubHosts[index2]
                    self.mcastSubtree[mflowId][dstHost] = dstSubHosts[:index2]
                    self.mcastSubtree[mflowId][host2] = dstSubHosts[index2+1:]
                    flow = assignNewFlow(dstHost, host2, self.flows, curFlow.flowSize/2, curTime, \
                                        mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)
                    
                    event = FlowArrivalEvent(curTime, flow)
                    eventQueue.push(event, event.time)
            else:
                #print 'McastFlow scatter finished'
                self.mcastHosts[mflowId] = {}
                self.mcastStage[mflowId] = Ring
                Hosts = [mcastFlow.srcHost]
                Hosts += mcastFlow.dstHosts.keys()
                for i in range(0, size):
                    host = Host()
                    host.hostId = Hosts[i]
                    host.rSubFlowIDs = [i+1]
                    host.sSubFlowIDs = [i+1]
                    host.subFlowSet = range(1, size+1)
                    self.mcastHosts[mflowId][Hosts[i]] = host
                    self.mcastFlows[mflowId].dstHosts[Hosts[i]] = False
                #print 'McastFlow transmission--ring'
                #print self.mcastHosts[mflowId]
                #print self.mcastHosts[mflowId].keys()
                for i in range(0, size):
                    srcHostId = self.mcastHosts[mflowId].keys()[i]
                    dstHostId = self.mcastHosts[mflowId].keys()[(i+1)%size]
                    srcHost = self.mcastHosts[mflowId].values()[i]
                    dstHost = self.mcastHosts[mflowId].values()[(i+1)%size]
                    if not (srcHost.sBusy or dstHost.rBusy):
                        srcHost.sBusy = True
                        dstHost.rBusy = True
                        subflowId = srcHost.sSubFlowIDs.pop(0)
                        flow = assignNewFlow(srcHostId, dstHostId, self.flows, sub_flowSize, curTime, \
                                            mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores, subflowId)
                        
                        event = FlowArrivalEvent(curTime, flow)
                        eventQueue.push(event, event.time)
        else: # Ring
            #print 'ring-------------'
            srcHostId = curFlow.startId
            dstHostId = curFlow.endId
            subflowId = curFlow.subflowId
            #print srcHostId, dstHostId
            
            index = self.mcastHosts[mflowId].keys().index(srcHostId)
            newDstHostId = self.mcastHosts[mflowId].keys()[(index+1)%len(self.mcastHosts[mflowId].keys())]
            srcHost = self.mcastHosts[mflowId].values()[index]
            dstHost = self.mcastHosts[mflowId].values()[(index+1)%len(self.mcastHosts[mflowId].values())]
            newDstHost = self.mcastHosts[mflowId].values()[(index+2)%len(self.mcastHosts[mflowId].values())]
            
            srcHost.sBusy = False
            dstHost.rBusy = False
            dstHost.rSubFlowIDs.append(subflowId)
            dstHost.sSubFlowIDs.append(subflowId)

            # receive all the subflows
            if set(dstHost.rSubFlowIDs) == set(dstHost.subFlowSet):
                mcastFlow.dstHosts[dstHostId] = True

            if all(mcastFlow.dstHosts.values()):
                mcastFlow.finishTime = curTime
                for flow in self.flows:
                    if flow.mcastFlowId == mflowId:
                        mcastFlow.energy += flow.energy
                print 'mcastFlow:%s transmission time %s energy %s'%(mflowId, mcastFlow.finishTime-mcastFlow.startTime, mcastFlow.energy)
            else:
                for i in xrange(0,len(self.mcastHosts[mflowId].keys())):
                    srcHostId = self.mcastHosts[mflowId].keys()[i]
                    dstHostId = self.mcastHosts[mflowId].keys()[(i+1)%len(self.mcastHosts[mflowId].keys())]
                    srcHost = self.mcastHosts[mflowId].values()[i]
                    dstHost = self.mcastHosts[mflowId].values()[(i+1)%len(self.mcastHosts[mflowId].values())] 
                    
                    if not (dstHost.sBusy or newDstHost.rBusy):
                        for subflowId in srcHost.sSubFlowIDs:
                            if subflowId not in dstHost.rSubFlowIDs:
                                flow = assignNewFlow(srcHostId, dstHostId, self.flows, sub_flowSize, curTime, \
                                            mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores, subflowId)
                                
                                event = FlowArrivalEvent(curTime, flow)
                                eventQueue.push(event, event.time)
                                srcHost.sSubFlowIDs.remove(subflowId)
                                